import torch
import torch.nn as nn


class Generator(nn.Module):
    """
    生成器
    """
    def __init__(self,noise_dim ,input_dim):
        super().__init__()
        self.input_dim = input_dim
        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat))
            layers.append(nn.ReLU(inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(noise_dim, 128, normalize=False),
            *block(128,256),
            *block(256,512),
            *block(512,1024),
            nn.Linear(1024,input_dim*input_dim),
            nn.Tanh()
        )

    def forward(self, x):
        img = self.model(x)
        img = img.view(-1,1, self.input_dim,self.input_dim)
        return img

class Discriminator(nn.Module):
    """
    判别器
    """
    def __init__(self,input_dim):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim*input_dim,512),
            nn.ReLU(inplace=True),
            nn.Linear(512,256),
            nn.ReLU(inplace=True),
            nn.Linear(256,1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.model(x)
        return x




if __name__ == '__main__':
    gen = Generator(100,28)
